import tushare as ts
import pandas as pd
import talib
import matplotlib.pyplot as plt
import os

# -------------------设置全局参数
params = {
    "股票代码": "000001.SZ",
#    "开始日期": pd.to_datetime("2023-1-1"),
#    "结束日期": pd.to_datetime("2025-1-1"),
    "开始日期": 20230101,
    "结束日期": 20250101,
    "初始资金": 100000,
    "快线天数": 20,
    "慢线天数": 60,
    "计划仓位": 200,
}

csv_fname="cached.csv"
if os.path.exists(csv_fname):
    df=pd.read_csv(csv_fname)
else:
    ts.set_token("0af3f3c3b4c2b1e34039449a85d25a13be5f75422c670b19d82ca553")
    pro=ts.pro_api()

    kparams = {'ts_code': '000001.SZ', 'startdate': 20230101, 'enddate': 20250101}
    df = pro.daily(ts_code = '000001.SZ',start_date = str(20230101),end_date=str(20250101))  # 从tushare中读取数据
    df.to_csv(csv_fname)

df_output = pd.DataFrame()

# -------------------初始化
def init():
    global df, df_output
    # factor计算
    df["sma_fast"] = talib.SMA(df["close"], timeperiod=params["快线天数"])
    df["sma_slow"] = talib.SMA(df["close"], timeperiod=params["慢线天数"])
    df["factor"] = df["sma_fast"] - df["sma_slow"]

    # 数据更新
    df = df.loc[df['trade_date'] > params["开始日期"]]
    df = df.loc[df['trade_date'] < params["结束日期"]]

    df_output = df[["trade_date", "close", "factor"]].copy()
    df_output = df_output[::-1]  # 按时间顺序

    # df_output=df_output.set_index(keys='trade_date')  #按时间顺序

    df_output[["资产", "现金", "持仓", "交易记录"]] = None
    df_output.iat[0, 3] = df_output.iat[0, 4] = params["初始资金"]
    df_output.iat[0, 5] = 0

# -------------------每个数据点执行函数
def on_bar(i):
    global df, df_output
    last_bar = df_output.iloc[i - 1]
    current_bar = df_output.iloc[i]
    #print(last_bar, current_bar)
    date = current_bar['trade_date']

    print(last_bar['trade_date'])
    pos = last_bar["持仓"]
    cash = last_bar["现金"]

    #    #["资产", "现金", "持仓", "交易记录"] 3,4,5,6

    # 开仓
    if not last_bar["持仓"]:
        if current_bar.factor > 0:
            cash = cash - current_bar.close * params["计划仓位"]
            pos = params["计划仓位"]
            df_output.iat[i, 6] = "buy"
            print(f"{date}, {current_bar.close}, buy")

    # 清仓
    if last_bar["持仓"]:
        if current_bar.factor < 0:
            cash = cash + current_bar.close * pos
            pos = 0
            df_output.iat[i, 6] = "sell"
            print(f"{date}, {current_bar.close}, sell")

        # 参数更新
    df_output.iat[i,4] = cash
    df_output.iat[i,5] = pos
    df_output.iat[i,3] = cash + current_bar['close'] * pos

# -------------------后处理
def output():
    profit_rate = df_output["资产"].iloc[-1] / df_output["资产"].iloc[0] - 1
    print(f"本次回测收益率：{profit_rate:.2%}")

    #显示中文
    plt.rcParams['font.sans-serif'] =  ['SimHei']  # 指定默认字体
    plt.rcParams['axes.unicode_minus'] = False  # 正确显示负号


    plt.subplot(2, 1, 1)
    plt.title(f'{params["股票代码"]}股价走势')
    df_output.close.plot()
    for trade_type in [("buy", "red"), ("sell", "blue")]:
        df_temp = df_output[df_output["交易记录"] == trade_type[0]]
        plt.scatter(
            df_temp.index,
            df_temp.close,
            color=trade_type[1],
            s=10,
            label=trade_type[0],
        )
    plt.legend()

    plt.subplot(2, 1, 2)
    plt.title("资产曲线")
    df_output["资产"].plot()
    plt.annotate(
        f"收益率：{profit_rate:.2%}",
        (df_output.index[-1], df_output["资产"].iloc[-1]),
    )
    plt.legend()
    plt.show()

# -------------------主函数
if __name__ == "__main__":
    init()
    for i in range(len(df.index)):
        if i == 0:
            continue
        on_bar(i)
    output()

